import pandas as pd
import pymysql
import os

# 下载表格
db_host = '127.0.0.1'
db_port = 3306
db_user = 'root'
db_password = 'root'
db_database = 'tfj'
db_charset = 'utf8mb4'

# 表名配置
table_name_q = 'goods_images_eval_o'
# table_name_v = 'goods_images_eval_v'

output_dir = 'data/split_results'
os.makedirs(output_dir, exist_ok=True)

def download_table_as_csv(table_name, db_params, output_filepath):
    """
    从数据库下载表并将其另存为 CSV 文件。
    """
    connection = None
    try:
        # 使用 pymysql 创建数据库连接
        connection = pymysql.connect(
            host=db_params['host'],
            port=db_params['port'],
            user=db_params['user'],
            password=db_params['password'],
            database=db_params['database'],
            charset=db_params['charset']
        )

        print(f"正在从 '{table_name}' 表下载数据...")
        # 构建查询语句并从数据库读取数据到 pandas DataFrame
        query = f"SELECT * FROM `{table_name}`"
        df = pd.read_sql_query(query, connection)

        # 将 DataFrame 保存为 CSV 文件
        df.to_csv(output_filepath, index=False, encoding='utf-8-sig')

        print(f"表 '{table_name}' 已成功保存到 '{output_filepath}'")

    except Exception as e:
        print(f"下载表 '{table_name}' 时出错: {e}")
    finally:
        if connection:
            connection.close()

# 数据库连接参数字典
db_params = {
    'user': db_user,
    'password': db_password,
    'host': db_host,
    'port': db_port,
    'database': db_database,
    'charset': db_charset
}

# 下载查询表
output_file_q_path = os.path.join(output_dir, f"{table_name_q}.csv")
download_table_as_csv(table_name_q, db_params, output_file_q_path)

# 下载数据库表
# output_file_v_path = os.path.join(output_dir, f"{table_name_v}.csv")
# download_table_as_csv(table_name_v, db_params, output_file_v_path)